import torch
import torchvision.transforms as transforms
from torchvision import models
from PIL import Image
import torch.nn as nn

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 定义图像预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # 将图像大小调整为模型期望的尺寸
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化
])

# 加载模型
model = models.vgg16()  # 初始化模型
num_features = model.classifier[6].in_features
model.classifier[6] = nn.Linear(num_features, 3)  # 根据分类任务调整最后一层
model.load_state_dict(torch.load('vgg16_best_model.pth'))  # 加载训练好的最佳模型权重
model = model.to(device)
model.eval()  # 设置为评估模式

# 加载图像
image_path_list = [
    r'D:\develop\PythonCode\python基础\附_项目实战\九_薄膜图片级别分类\data\real_data\level_1\1.tif',
    r'D:\develop\PythonCode\python基础\附_项目实战\九_薄膜图片级别分类\data\real_data\level_2\1.tif',
    r'D:\develop\PythonCode\python基础\附_项目实战\九_薄膜图片级别分类\data\real_data\level_3\1.tif'

                   ]  # 替换为你的图片路径
for image_path in image_path_list:
    image = Image.open(image_path).convert('RGB')  # 打开图片
    image = transform(image)  # 应用预处理
    image = image.unsqueeze(0)  # 增加批次维度
    image = image.to(device)  # 将图像移至设备

    # 预测
    with torch.no_grad():
        outputs = model(image)
        _, predicted = torch.max(outputs, 1)
        predicted = predicted.cpu().numpy()

    # 显示预测结果
    class_names = ['最优级-1级', '中间级-2级', '差等级-3级']  # 类别名
    print("Predicted Class:", class_names[predicted[0]])
